import streamlit as st
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
st.set_page_config(page_title="nn_forward_propagation", page_icon="🧠")

# -----------------------
# 定义简单神经网络
# -----------------------
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(2, 4)  # 输入2 -> 隐藏层4
        self.fc2 = nn.Linear(4, 1)  # 隐藏层4 -> 输出1

    def forward(self, x):
        hidden = torch.relu(self.fc1(x))
        output = self.fc2(hidden)
        return hidden, output

# -----------------------
# Streamlit 前端界面
# -----------------------

st.title("🧠 神经网络前向传播演示")
st.caption("作者：何双新 ｜ 功能：神经网络前向传播演示")
st.write("这里是神经网络前向传播的小工具展示页。")
st.markdown("""
输入任意2个特征值，我们实时展示：
- 隐藏层激活值
- 最终输出
- 可视化神经元输出变化
""")

# 输入数据
st.subheader("🔢 输入数据")
input1 = st.number_input("输入特征1", value=1.0)
input2 = st.number_input("输入特征2", value=2.0)

if st.button("执行前向传播并绘图"):
    # 1. 初始化模型
    model = SimpleNN()
    model.eval()

    # 2. 组织输入
    input_tensor = torch.tensor([[input1, input2]], dtype=torch.float32)

    # 3. 前向传播
    with torch.no_grad():
        hidden, output = model(input_tensor)

    hidden_np = hidden.numpy().flatten()
    output_np = output.numpy().flatten()

    # 4. 显示数值
    st.subheader("🧠 隐藏层输出（激活值）")
    st.write(hidden_np)

    st.subheader("🎯 最终输出")
    st.success(f"输出值: {output_np[0]:.4f}")

    # 5. 绘制神经元输出图
    st.subheader("📈 神经元激活图")

    fig, ax = plt.subplots(figsize=(8, 4))
    bars = ax.bar(range(1, len(hidden_np) + 2), np.append(hidden_np, output_np), color=['skyblue'] * 4 + ['salmon'])
    ax.set_xticks(range(1, len(hidden_np) + 2))
    ax.set_xticklabels([f"Hidden {i+1}" for i in range(len(hidden_np))] + ["Output"])
    ax.set_ylabel("激活值")
    ax.set_ylim(0, max(1.0, hidden_np.max() + 0.5))

    # 添加数值标签
    for bar in bars:
        height = bar.get_height()
        ax.annotate(f'{height:.2f}',
                    xy=(bar.get_x() + bar.get_width() / 2, height),
                    xytext=(0, 3),
                    textcoords="offset points",
                    ha='center', va='bottom', fontsize=9)

    st.pyplot(fig)

st.markdown("---")
st.caption("安徽智加数字科技有限公司 · 技术学习组出品 🚀")
